import pandas as pd
import numpy as np
import statsmodels.api as sm
import warnings
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")

data_train = pd.read_csv('data_train.txt', index_col=[0], header=0)
data_predict = pd.read_csv('data_predict.txt', index_col=[0], header=0)

# 补充是否为夜晚（2点到5点）的二元变量
data_train['night'] = data_train['time_hour'].apply(lambda x: 1 if (x > 1) & (x < 6) else 0)
data_predict['night'] = data_predict['time_hour'].apply(lambda x: 1 if (x > 1) & (x < 6) else 0)

train_subset = data_train.loc[:,
               ['repost', 'comments', 'likes', 'tfidf', 'number_in_train', 'forward_max', 'comment_max', 'like_max',
                'forward_mean',
                'comment_mean', 'like_mean', 'time_weekend', 'panduan', 'length_all', 'length_chinese', 'sharing',
                'book',
                'mention', 'emoji', 'video', 'http', 'title', 'hotwords', 'keywords', 'is_noise', 'stock', 'night',
                'lottery']]

# 将互动量为0的数据标记为is_zero
train_subset['is_zero'] = 0
train_subset['is_zero'][(train_subset['comments'] + train_subset['repost'] + train_subset['likes']) == 0] = 1

# 设置模型的自变量和因变量
y_train = train_subset['is_zero']
x_train = train_subset[
    ['tfidf', 'number_in_train', 'forward_max', 'forward_mean', 'time_weekend', 'length_all', 'sharing', 'book',
     'mention', 'emoji', 'video', 'http', 'title', 'hotwords', 'keywords', 'stock', 'is_noise', 'night', 'lottery']]
x_train = sm.add_constant(x_train)

# 拟合logit模型
lr = sm.Logit(y_train, x_train)
result = lr.fit()
print(result.summary())

y_hat = result.predict(x_train)
compare = pd.DataFrame({'predict': y_hat, 'actual': train_subset['is_zero']})

# 设置不同的阈值，计算TPR和FPR
criterion = np.arange(0.3, 1, 0.01)
num = len(criterion)
TPR = np.zeros(num)
FPR = np.zeros(num)
true_index = compare['actual'] == 1
false_index = compare['actual'] == 0
for ii in range(num):
    compare['y_hat'] = compare['predict'] > criterion[ii]
    TPR[ii] = sum(compare['y_hat'][true_index]) / sum(true_index)
    FPR[ii] = sum(compare['y_hat'][false_index]) / sum(false_index)

# 画出ROC曲线
fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(FPR, TPR)
ax.set_title('ROC Curve')
ax.set_xlabel('False Positive Ratio')
ax.set_ylabel('True Positive Ratio')
plt.show()

# 确认最优的阈值0.67
distance = (1 - TPR) ** 2 + FPR ** 2
optimal_index = np.argmin(distance)
threshold = criterion[optimal_index]
threshold

# 根据最优阈值设立logit的判断结果，1为0互动量，0为有互动量数据
data_train['logit'] = compare['predict'].apply(lambda x: 1 if x > threshold else 0)
predict_subset = data_predict.loc[:,
                 ['tfidf', 'number_in_train', 'forward_max', 'comment_max', 'like_max', 'forward_mean',
                  'comment_mean', 'like_mean', 'time_weekend', 'panduan', 'length_all', 'length_chinese', 'sharing',
                  'book',
                  'mention', 'emoji', 'video', 'http', 'title', 'hotwords', 'keywords', 'is_noise', 'stock', 'night',
                  'lottery']]
x_predict = predict_subset[
    ['tfidf', 'number_in_train', 'forward_max', 'forward_mean', 'time_weekend', 'length_all', 'sharing', 'book',
     'mention', 'emoji', 'video', 'http', 'title', 'hotwords', 'keywords', 'stock', 'is_noise', 'night', 'lottery']]
x_predict = sm.add_constant(x_predict)
y_hat_predict = result.predict(x_predict)
data_predict['logit'] = (y_hat_predict > threshold).astype(int)

data_train.to_csv('data_train.txt', header=True)
data_predict.to_csv('data_predict.txt', header=True)
